import io
import os
from datetime import datetime, timedelta

import pandas as pd
from scapy.packet import Raw
from scapy.utils import rdpcap
from xbase_util.add_column_util import handle_dns
from xbase_util.common_util import build_es_expression, extract_session_fields, process_origin_pos
from xbase_util.handle_features_util import handle_uri, handle_ua
from xbase_util.pcap_util import fix_pos, group_numbers, get_file_and_read_pos
from global_ver import es_req, global_geo_util, base_config
from src.bean.TrafficDetailBean import TrafficDetailBean
from src.util.package import PacketUtil


def get_traffic_detail(bean: TrafficDetailBean, common_logger):
    if "-" in bean.start_time:
        start_time = datetime.strptime(bean.start_time, '%Y-%m-%d %H:%M:%S')
    elif "/" in bean.start_time:
        start_time = datetime.strptime(bean.start_time, '%Y/%m/%d %H:%M:%S')
    else:
        return {"code": 501, "message": "开始时间解析失败，start_time不是%Y-%m-%d %H:%M:%S或%Y/%m/%d %H:%M:%S",
                "data": None}
    protocols = bean.protocols
    protocols = f"{protocols}".replace("[", "").replace("]", "").replace("\"", "").replace("'", "").split(",")
    protocols=[item for item in protocols if item.strip()!=""]
    common_logger.info(f"[获取数据详情]获取协议：{protocols}")
    exps = [f"port.src=={bean.spt}", f"port.dst=={bean.dpt}"] + [f"protocols=={protocol}" for protocol in
                                                                 protocols]
    if bean.sip != "":
        exps.append(f"ip.src=={bean.sip}")
    if bean.dip != "":
        exps.append(f"ip.dst=={bean.dip}")
    common_logger.info(f"[获取数据详情]表达式：{'&&'.join(exps)}")
    common_logger.info(f"[获取数据详情]时间:{start_time}")
    es_exp = build_es_expression(1, "&&".join(exps), start_time - timedelta(seconds=5),
                                 start_time + timedelta(seconds=5), bounded_type="last")
    common_logger.info(f"[获取数据详情]已创建:{es_exp}")
    res = es_req.search(es_exp, "arkime_sessions3*")
    common_logger.info(f"[获取数据详情]搜索结果:{res.text}")
    res_data = res.json()['hits']['hits']
    if len(res_data) == 0:
        common_logger.info(f"[获取数据详情]数据为空:{es_exp}")
        return {"code": 200, "message": "数据为空，请尝试其他时间范围", "data": None}
    # res_data = res_data[0]['_source']
    data_list = extract_session_fields(res_data, global_geo_util)
    common_logger.info(f"[获取数据详情]数据量：{len(data_list)}")
    df = handle_uri(pd.DataFrame(data_list), use_tqdm=False)
    df = handle_ua(df, use_tqdm=False)
    df = handle_dns(df, isDataFrame=True, use_tqdm=False)
    session = df.to_dict(orient='records')[0]
    common_logger.info(f"[获取数据详情]处理完毕：{len(session)}")
    originPos = session['packetPos']
    pos = process_origin_pos(originPos)
    prefix = base_config['path_prefix']
    hits = es_req.search_file(f"{session['node']}-{abs(pos[0])}").json()['hits']['hits']
    common_logger.info(f"[获取数据详情]搜索pcap文件：{hits}")
    file = None
    if len(hits) > 0:
        file = hits[0]['_source']
        if prefix is None:
            return file
        prefix_res = prefix
        if not prefix.endswith('/'):
            prefix_res = f"{prefix}/"
        origin_path = file['name']
        basename = os.path.basename(origin_path)
        result_path = f"{prefix_res}{basename}"
        file['name'] = result_path
    common_logger.info(f"[获取数据详情]文件：{file}")
    if file is None:
        pcap_flow_text = ""
    else:
        fix_pos(pos, file['packetPosEncoding'])
        pos_list = group_numbers(pos)[0]
        pos_list.pop(0)
        stream = get_file_and_read_pos(session['id'], file, pos_list)
        if stream is None:
            common_logger.info("[获取数据详情]文件信息获取失败")
            return [PacketUtil.edit_session_and_save(session=session)]
        else:
            packets = rdpcap(io.BytesIO(stream))
            streams = b""
            for pkt in packets:
                if Raw in pkt:
                    streams += pkt[Raw].load
            pcap_flow_text = streams.decode('ascii', errors='ignore')
    session['pcap_flow_text'] = pcap_flow_text
    common_logger.info("[获取数据详情]获取成功")
    return {"data": session, "message": "获取成功", "code": 200}
